#include "arcball.h"
#include "constant.h"
#include "gn_solver.h"
#include <helper_cuda.h>
#include <stdlib.h>
//#include <GL/glut.h>
#include <GL/freeglut.h>
#include <fstream>

#include <innoreal/utils/innoreal_helper.hpp>
#include <innoreal/utils/innoreal_timer.hpp>
#include <math_functions.h>

float3 bound_length, min_corner;
float max_length;
static float scal = 1.0, mag = 1.3;
// arcball parameters
static float aspect_ratio = 1.0f;
static int width, height;
float scale = 65.0f;
const vec eye(0.0f, 0.0f, -15.0f);
vec centre(0.0f, 0.0f, 0.0f);
bool g_running = true;
const vec up(0.0f, 1.0f, 0.0f);
char global_meshdir[256];

thrust::host_vector<int4> h_face_src, h_face_tar;
thrust::host_vector<float4> h_ver_src, h_ver_tar, h_ver_deformed, h_nor_src, h_nor_tar, h_nor_deformed, h_node_src, h_node_deformed;
thrust::host_vector<int> h_node_rela_ind_src, h_node_rela_ind_deformed, h_corr_pair_list;

bool node_flag_src = false, err_flag = true, src_flag = true, tar_flag = true, ori_flag = false;
GNSolver *gnSolver_ptr = new GNSolver;
InputData *input_data_ptr = new InputData;

int source_seri, target_seri;
bool device_flag = false;

float4 corr_pair_left, corr_pair_right;
bool corr_pair_flag = false;
int trans = 0.01;

char fn_deformed[100];
int iter_count = 0;

innoreal::InnoRealTimer timer;

static void reset_view(int w, int h) {
  width = w;
  height = h;
  aspect_ratio = (float) width / (float) height;
  glViewport(0, 0, width, height);
  glMatrixMode(GL_PROJECTION);
  glLoadIdentity();
  gluPerspective(10.0f, aspect_ratio, 1.0f, 50.f);
  gluLookAt(
      eye.x, eye.y, eye.z,
      centre.x, centre.y, centre.z,
      up.x, up.y, up.z);
  // set up the arcball using the current projection matrix
  //    arcball_setzoom(0.002, eye, up);
  arcball_setzoom(1.0f, eye, up);

  glMatrixMode(GL_MODELVIEW);
  glLoadIdentity();
}

void InitIndata(const char *meshdir, int source_id, int target_id) {
  source_seri = source_id;
  target_seri = target_id;

  char file_src[1024], file_tar[1024];

  sprintf(file_src, "%s/mesh%d.obj", meshdir, source_id);
  sprintf(file_tar, "%s/mesh%d.obj", meshdir, target_id);
  
	// sprintf(file_src, "%s/%d_cleaner.obj", meshdir, source_id);
	// sprintf(file_tar, "%s/%d_cleaner.obj", meshdir, target_id);

  timer.TimeStart();
  input_data_ptr->Init(file_src, file_tar);
  timer.TimeEnd();
	INNOREALI("time of prepare data: %f ms\n", timer.TimeGap_in_ms());
  input_data_ptr->source_.ComputeBoundBox();
  bound_length.x = input_data_ptr->source_.h_bound_box_[1] - input_data_ptr->source_.h_bound_box_[0];
  bound_length.y = input_data_ptr->source_.h_bound_box_[3] - input_data_ptr->source_.h_bound_box_[2];
  bound_length.z = input_data_ptr->source_.h_bound_box_[5] - input_data_ptr->source_.h_bound_box_[4];
  max_length = max(max(bound_length.x, bound_length.y), bound_length.z);
  centre.x = bound_length.x / max_length / 2.f;
  centre.y = bound_length.y / max_length / 2.f;
  centre.z = bound_length.z / max_length / 2.f;
  min_corner.x = input_data_ptr->source_.h_bound_box_[0];
  min_corner.y = input_data_ptr->source_.h_bound_box_[2];
  min_corner.z = input_data_ptr->source_.h_bound_box_[4];

  h_ver_src = input_data_ptr->source_.d_vertex_;
  h_ver_tar = input_data_ptr->target_.d_vertex_;
  h_ver_deformed = input_data_ptr->deformed_.d_vertex_;
  h_face_src = input_data_ptr->source_.d_face_;
  h_face_tar = input_data_ptr->target_.d_face_;
  h_nor_src = input_data_ptr->source_.d_normal_;
  h_nor_deformed = input_data_ptr->deformed_.d_normal_;
  h_nor_tar = input_data_ptr->target_.d_normal_;
  h_node_src = input_data_ptr->source_.d_node_;
  h_node_deformed = input_data_ptr->deformed_.d_node_;
  h_node_rela_ind_src = input_data_ptr->source_.node_rela_index_;
  h_node_rela_ind_deformed = input_data_ptr->deformed_.node_rela_index_;
  h_corr_pair_list = input_data_ptr->d_corr_index_;

  SolverPara gnsolver_param(true, 10.0f, false, 10000.0f, true, 3.0f, true, 50.0f);
  //    SolverPara gnsolver_param(true, 1000.0f, true, 300.0f, true, 0.1f, 1, 240000.0f);

  timer.TimeStart();
  gnSolver_ptr->InitSovler(input_data_ptr, gnsolver_param);
  timer.TimeEnd();
  INNOREALI("time of init gnSolver: %f ms\n", timer.TimeGap_in_ms());
  device_flag = true;
}

void renderScene(void) {
  if (!device_flag) {
    return;
  }
  if (!g_running) {
    return;
  }
  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

  glMatrixMode(GL_MODELVIEW);
  glLoadIdentity();
  glPushMatrix();
  glTranslatef(centre.x, centre.y, centre.z);
  arcball_rotate();
  glScalef(scal, scal, scal);

  if (1) {
    float4 vertex, normal;
    int4 face;
    if (src_flag) {
      h_ver_deformed = input_data_ptr->deformed_.d_vertex_;
      h_face_src = input_data_ptr->source_.d_face_;
      h_nor_deformed = input_data_ptr->deformed_.d_normal_;
      glColor3f(1, 0, 0);
      //	glPointSize(4);
      glBegin(GL_TRIANGLES);
      for (int i = 0; i < input_data_ptr->GetSrcFaceNum(); i++) {
        face = h_face_src[i];
        if (face.x < 0 || face.y < 0 || face.z < 0) {
          continue;
        }
        vertex = h_ver_deformed[face.x];
        normal = h_nor_deformed[face.x];
        glVertex3f((vertex.x - min_corner.x) / max_length, (vertex.y - min_corner.y) / max_length, (vertex.z - min_corner.z) / max_length);
        glNormal3f(-normal.x, -normal.y, -normal.z);

        vertex = h_ver_deformed[face.y];
        normal = h_nor_deformed[face.y];
        glVertex3f((vertex.x - min_corner.x) / max_length, (vertex.y - min_corner.y) / max_length, (vertex.z - min_corner.z) / max_length);
        glNormal3f(-normal.x, -normal.y, -normal.z);

        vertex = h_ver_deformed[face.z];
        normal = h_nor_deformed[face.z];
        glVertex3f((vertex.x - min_corner.x) / max_length, (vertex.y - min_corner.y) / max_length, (vertex.z - min_corner.z) / max_length);
        glNormal3f(-normal.x, -normal.y, -normal.z);
      }
      glEnd();
    }

    if (tar_flag) {
      h_ver_tar = input_data_ptr->target_.d_vertex_;
      h_face_tar = input_data_ptr->target_.d_face_;
      h_nor_tar = input_data_ptr->target_.d_normal_;
      glColor3f(0, 1, 0);
      //	glPointSize(4);
      glBegin(GL_TRIANGLES);
      for (int i = 0; i < input_data_ptr->GetTarFaceNum(); i++) {
        face = h_face_tar[i];
        if (face.x < 0 || face.y < 0 || face.z < 0) {
          continue;
        }
        vertex = h_ver_tar[face.x];
        normal = h_nor_tar[face.x];
        glVertex3f((vertex.x - min_corner.x) / max_length, (vertex.y - min_corner.y) / max_length, (vertex.z - min_corner.z) / max_length);
        glNormal3f(-normal.x, -normal.y, -normal.z);

        vertex = h_ver_tar[face.y];
        normal = h_nor_tar[face.y];
        glVertex3f((vertex.x - min_corner.x) / max_length, (vertex.y - min_corner.y) / max_length, (vertex.z - min_corner.z) / max_length);
        glNormal3f(-normal.x, -normal.y, -normal.z);

        vertex = h_ver_tar[face.z];
        normal = h_nor_tar[face.z];
        glVertex3f((vertex.x - min_corner.x) / max_length, (vertex.y - min_corner.y) / max_length, (vertex.z - min_corner.z) / max_length);
        glNormal3f(-normal.x, -normal.y, -normal.z);
      }
      glEnd();
    }

    if (corr_pair_flag) {
      glColor3f(0, 1, 1);
      glPointSize(3);
      glBegin(GL_POINT);
      glVertex3f((corr_pair_left.x - min_corner.x) / max_length, (corr_pair_left.y - min_corner.y) / max_length, (corr_pair_left.z - min_corner.z) / max_length);
      glEnd();

      glColor3f(1, 1, 0);
      glPointSize(3);
      glBegin(GL_POINT);
      glVertex3f((corr_pair_right.x - min_corner.x) / max_length, (corr_pair_right.y - min_corner.y) / max_length, (corr_pair_right.z - min_corner.z) / max_length);
      glEnd();

      glColor3f(1, 1, 1);
      glBegin(GL_LINES);
      glVertex3f((corr_pair_left.x - min_corner.x) / max_length, (corr_pair_left.y - min_corner.y) / max_length, (corr_pair_left.z - min_corner.z) / max_length);
      glVertex3f((corr_pair_right.x - min_corner.x) / max_length, (corr_pair_right.y - min_corner.y) / max_length, (corr_pair_right.z - min_corner.z) / max_length);
      glEnd();
    }
  }

  if (node_flag_src) {
    h_node_src = input_data_ptr->source_.d_node_;
    h_node_rela_ind_src = input_data_ptr->source_.node_rela_index_;
    float4 temp_node_l, temp_node_r;
    glColor3f(1, 0, 0);
    glPointSize(3);
    glBegin(GL_LINES);
    //       glBegin(GL_POINTS);
    for (int i = 0; i < h_node_src.size(); i++) {
      temp_node_l = h_node_src[i];
      for (int j = 1; j < MAX_NEAR_NODE_NUM_NODE; j++) {
        temp_node_r = h_node_src[h_node_rela_ind_src[i * MAX_NEAR_NODE_NUM_NODE + j]];
        glVertex3f((temp_node_l.x - min_corner.x) / max_length, (temp_node_l.y - min_corner.y) / max_length, (temp_node_l.z - min_corner.z) / max_length);
        glVertex3f((temp_node_r.x - min_corner.x) / max_length, (temp_node_r.y - min_corner.y) / max_length, (temp_node_r.z - min_corner.z) / max_length);
      }
    }
    glEnd();
  }
  glPopMatrix();

  // ����������
  glutSwapBuffers();

  glFlush();
}

void processNormalKeys(unsigned char key, int x, int y) {

  if (key == 27) {
    glutLeaveMainLoop();
    g_running = false;
  }
  if (key == 45) { // - zoom out
    scal -= 0.05f;
  } else if (key == 61) { // + zoom in
    scal += 0.05f;
  } else if (key == 48) { // 0
    src_flag ^= true;
  } else if (key == 49) { // 1
    tar_flag ^= true;
  } else if (key == 50) { // 2
    node_flag_src ^= true;
  }
  //	 cout << "Key:" << int(key) << endl;
}

void nonrigidICP(int flag);

void DrawCorrespondingPair() {
  int id_l = 0, id_r = 0;
  std::cout << "input the ID of vertex in source: ";
  std::cin >> id_l;
  if (id_l < 0 || id_l > input_data_ptr->GetSrcVerNum() - 1) {
    printf("invalid number, vertex ID should be one of[0, %d]!\n", input_data_ptr->GetSrcVerNum() - 1);
    corr_pair_flag = false;
    return;
  }
  h_ver_deformed = input_data_ptr->deformed_.d_vertex_;
  corr_pair_left = h_ver_deformed[id_l];
  id_r = h_corr_pair_list[id_l];
  if (id_r > -1) {
    corr_pair_right = h_ver_tar[id_r];
    corr_pair_flag = true;
  } else {
    printf("invalid corresponding pair,left:%d\n", id_l);
  }
}

void RegisteMesh() {
  int iter_num = 1;
  for (int i = 0; i < iter_num; i++) {
    timer.TimeStart();
    nonrigidICP(1);
    timer.TimeEnd();
    INNOREALI("time of %d iters: %f ms\n", i + 1, timer.TimeGap_in_ms());
    iter_count++;
    timer.TimeStart();
    renderScene();
    timer.TimeEnd();
    INNOREALI("time of rendering: %.6f ms \n", timer.TimeGap_in_ms());
  }
}

void UpdateTargetMesh() {
  char update_tar_fn[100];
  target_seri++;
  //  sprintf(update_tar_fn, "/media/vradmin/0006CA8D00008838/data/original_sfm_data/result/%d.obj", target_seri);
  sprintf(update_tar_fn, "%s/%d.obj", global_meshdir, target_seri);
  //  sprintf(update_tar_fn, "/home/vradmin/workspace/xjm/xjm_pointclouds/xjm_pointclouds/%06d_t___new___bew.obj", target_seri);

  input_data_ptr->UpdateTarget(update_tar_fn);
  gnSolver_ptr->counter_iter_ = 0;
}

void SaveResult() {
  char target_mesh_fn[100];
  char deformed_node_fn[100];
  char deformed_mesh_fn[100];
  char tforms_fn[100];
  sprintf(target_mesh_fn, "save/source_mesh_%d.obj", target_seri);
  sprintf(deformed_node_fn, "save/deformed_node_%d_%d.obj", source_seri, target_seri);
  sprintf(tforms_fn, "save/tforms_%d_%d.txt", source_seri, target_seri);
  sprintf(deformed_mesh_fn, "save/deformed_mesh_%d_%d.obj", source_seri, target_seri);
  //    input_data_ptr->source_.SaveNodeToFile(source_node_fn, 0);
  //    input_data_ptr->deformed_.SaveNodeToFile(deformed_node_fn, 1);
  //    gnSolver_ptr->SaveTformsToFile(tforms_fn);
  //    input_data_ptr->deformed_.SaveMeshToFile(deformed_mesh_fn, 1);
  input_data_ptr->target_.SaveMeshToFile(target_mesh_fn, 1);
}

void SaveSrcAndTarMesh() {
  char source_mesh_fn[100];
  char target_mesh_fn[100];
  char deformed_mesh_fn[100];
  char node_graph_fn[100];
  sprintf(source_mesh_fn, "save/source_mesh_%d.obj", source_seri);
  sprintf(target_mesh_fn, "save/target_mesh_%d.obj", target_seri);
  sprintf(deformed_mesh_fn, "save/deformed_mesh_%d.obj", target_seri);
  sprintf(node_graph_fn, "save/node_graph_%d.obj", source_seri);
  input_data_ptr->source_.SaveMeshToFile(source_mesh_fn, 2);
  input_data_ptr->target_.SaveMeshToFile(target_mesh_fn, 3);
  input_data_ptr->deformed_.SaveMeshToFile(deformed_mesh_fn, 1);
  input_data_ptr->source_.SaveNodeGraphWithLineToFile(node_graph_fn);
  //  input_data_ptr->source_.
}

void processSpecialKeys(int key, int x, int y) {

  switch (key) {
    case GLUT_KEY_F5:
      for (int iter = 0; iter < 1; ++iter) {
        //        char deformed_mesh_fn[100];
        //        sprintf(deformed_mesh_fn, "save/deformed_mesh_%d_%d_%d.obj", source_seri, target_seri, gnSolver_ptr->counter_iter_);
        //        input_data_ptr->deformed_.SaveMeshToFile(deformed_mesh_fn, 1);
        RegisteMesh();
      }
#if 0
    {
      char target_mesh_fn[100];
      sprintf(target_mesh_fn, "save/target_mesh_%d.obj", target_seri);
      input_data_ptr->target_.SaveMeshToFile(target_mesh_fn, 3);
    }
#endif
      break;
    case GLUT_KEY_F6:char target_mesh_fn[100];
      sprintf(target_mesh_fn, "save/target_mesh_%d.obj", target_seri);
      input_data_ptr->target_.SaveMeshToFile(target_mesh_fn, 3);
      break;
    case GLUT_KEY_F7:UpdateTargetMesh();
      break;
    case GLUT_KEY_F8:DrawCorrespondingPair();
      break;
    case GLUT_KEY_F9:
      for (int i = 0; i < 15; i++) {
        RegisteMesh();
        //            SaveResult();
        //        SaveSrcAndTarMesh();
        UpdateTargetMesh();
      }
      break;
    case GLUT_KEY_F10:input_data_ptr->deformed_.SaveMeshToFile("save/deformed.obj", 1);
      //      input_data_ptr->deformed_.SaveNodeToFile("save/node_deformed.obj", 1);
      input_data_ptr->source_.SaveMeshToFile("save/source.obj", 2);
      input_data_ptr->source_.SaveNodeToFile("save/node_source.obj", 1);
      input_data_ptr->target_.SaveMeshToFile("save/target.obj", 3);
      input_data_ptr->SaveCorrToFile("save/corr.obj");
      input_data_ptr->source_.SavePointCloudAndNormalToFile("save/source_normal.ply");
      input_data_ptr->target_.SavePointCloudAndNormalToFile("save/target_normal.ply");
            input_data_ptr->source_.SaveNodeGraphWithLineToFile("save/node_graph_source.obj");
      //      input_data_ptr->source_.SavePointCloudAndNormalToFile("save/point_cloud_source.ply");
      //      input_data_ptr->target_.SavePointCloudAndNormalToFile("save/point_cloud_target.ply");
      //      gnSolver_ptr->SaveTformsToFile("save/vars.txt");

      break;
    case GLUT_KEY_F11:char source_node_fn[100];
      char deformed_node_fn[100];
      char deformed_mesh_fn[100];
      char tforms_fn[100];
      sprintf(source_node_fn, "save/source_node_%d_%d.obj", source_seri, target_seri);
      sprintf(deformed_node_fn, "save/deformed_node_%d_%d.obj", source_seri, target_seri);
      sprintf(tforms_fn, "save/tforms_%d_%d.txt", source_seri, target_seri);
      sprintf(deformed_mesh_fn, "save/deformed_mesh_%d_%d.obj", source_seri, target_seri);
      input_data_ptr->source_.SaveNodeToFile(source_node_fn, 0);
      input_data_ptr->deformed_.SaveNodeToFile(deformed_node_fn, 1);
      gnSolver_ptr->SaveTformsToFile(tforms_fn);
      input_data_ptr->deformed_.SaveMeshToFile(deformed_mesh_fn, 1);
      break;
  }
}

static void resize(int w, int h) {
  reset_view(w, h);
}

static void click_scene(int x, int y) {
  int invert_y = (height - y) - 1; // OpenGL viewport coordinates are Cartesian
  arcball_start(x, invert_y);
}

static void drag_scene(int x, int y) {
  int invert_y = (height - y) - 1;
  arcball_move(x, invert_y);
}

static void mouse_button(int button, int state, int x, int y) {
  if (state == GLUT_DOWN) {
    click_scene(x, y);
  }
}

static void mouse_motion(int x, int y) {
  // glutMotionFunc only called when a mouse button is held down
  drag_scene(x, y);
}

int main(int argc, char **argv) {
	memset(global_meshdir, 0, sizeof(global_meshdir));
	sprintf(global_meshdir, "%s", "/data/workspace/Offline4D/data/168/mesh");
	// sprintf(global_meshdir, "%s", "/data1/Dataset/meshdata/original_sfm_data/mesh");

  InitIndata(global_meshdir, atoi(argv[1]), atoi(argv[2]));

  arcball_reset();
  glutInit(&argc, argv);
  glutInitDisplayMode(GLUT_DEPTH | GLUT_DOUBLE | GLUT_RGBA);
  glutInitWindowPosition(500, 100);
  glutInitWindowSize(640, 640);
  glutCreateWindow("3D Tech- GLUT Tutorial");
  glutDisplayFunc(renderScene);
  glutKeyboardFunc(processNormalKeys);
  glutIdleFunc(renderScene);
  glutMouseFunc(mouse_button);
  glutMotionFunc(mouse_motion);
  glutSpecialFunc(processSpecialKeys);
  glutReshapeFunc(resize);
  glutSetOption(GLUT_ACTION_ON_WINDOW_CLOSE, GLUT_ACTION_CONTINUE_EXECUTION);
  GLfloat globel_ambient[] = {0.0, 0.0, 0.0, 1.0};
  // global light on
  glLightModelfv(GL_LIGHT_MODEL_AMBIENT, globel_ambient);

  glEnable(GL_LIGHTING);
  glEnable(GL_LIGHT0);

  glEnable(GL_DEPTH_TEST);

  glEnable(GL_COLOR_MATERIAL);

  glutMainLoop();

  delete input_data_ptr;
  delete gnSolver_ptr;
  return 0;
}

void nonrigidICP(int flag) {
  gnSolver_ptr->Run(flag);
}
